import os
import os.path as osp
from glob import glob
import numpy as np
import cv2 as cv
import re
from tqdm import tqdm

def extract_number(filename):
    # 从文件名中提取数字部分
    match = re.search(r'\d+', filename)
    if match:
        return int(match.group())
    return 0


if __name__=="__main__":
    base_path = '/home/dataset/2024-06-05-ARMOT/train'
    image_paths = os.listdir(base_path)
    image_paths.sort()
    for image_path in tqdm(image_paths):
        img_dir = osp.join(base_path, image_path, 'img1')
        imgs = os.listdir(img_dir)
        imgs = sorted(imgs,key=extract_number)
        # print(imgs)
        gt_path = osp.join(base_path,image_path,'gt','gt.txt')
        gt_ori = np.loadtxt(gt_path,dtype=int,delimiter=',')
        len_gt = len(gt_ori)
        mot_cat = [[1,1,1] for i in range(len_gt)]
        avl_colors = {
            165: [183, 10, 186],
            166: [165, 28, 61],
            167: [144, 238, 194],
            168: [12, 158, 41],
            169: [76, 110, 234],
            1: [153, 108, 6],
            3: [89, 121, 72],
            7: [115, 176, 195],
            8: [161, 171, 27],
            9: [135, 169, 180],
            # 165: [183, 10, 186],
            10: [29, 26, 199],
            12: [242, 107, 146],
            13: [156, 198, 23],
            11: [102, 16, 239],  ###########SDrone2 before
            170: [150, 9, 121],
            171: [142, 1, 246]
        }
        segs = avl_colors.keys()
        segs = list(segs)
        gt_list = []
        for fid, img in enumerate(imgs):
            img_path = osp.join(img_dir,img)
            # print(img_path)
            show_image = cv.imread(img_path)
            mask = gt_ori[:,0]==fid+1#fid与image名字对应
            annos = gt_ori[mask]
            for anno in annos:
                # fid,x,y,w,h,seg =anno
                fid,id,x,y,w,h,_,_,_ = anno[:9]
                bottom_right_x = int(x + w )
                bottom_right_y = int(y + h )
                pt_1 = (x, y)
                pt_2 = (bottom_right_x,bottom_right_y)

                cv.rectangle(show_image,pt_1,pt_2,(0,0,255),2)

            anno_path = osp.join(img_dir[:-4],'anno')
            # print(anno_path)

            os.makedirs(anno_path,exist_ok=True)
            #anno_image_name = osp.join(anno_path,img)
            #cv.imwrite(anno_image_name, show_image)

            cv.imshow('test', show_image)
            q = cv.waitKey(50)
            if q == ord('q'):
                break
    cv.destroyAllWindows()
